# Copyright 2023 The jax3d Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Precompyted Glasbey Colormap."""

import functools

import jax.numpy as jnp
from jax3d.projects.nesf.utils.typing import i32


@functools.lru_cache()
def get_color_map() -> i32["255 3"]:
  """Returns the colormap.

  The colormap was precomputed using Glasbey et al. algorythm (Colour Displays
  for Categorical Images, 2017) to generate maximally distinct colors.

  It was generated using https://github.com/taketwo/glasbey:

  ```python
  gb = glasbey.Glasbey(
      base_palette=[(0, 0, 0), (228, 26, 28), (55, 126, 184), (77, 175, 74)],
      no_black=True,
  )
  palette = gb.generate_palette(size=256)
  gb.save_palette(palette, 'colormap.csv')
  ```

  Returns:
    colormap: A `jnnp.array(shape=(256, 3), dtype=np.uint8)` representing the
      mapping id -> color.
  """
  return jnp.asarray([
      0, 0, 0,
      228, 26, 28,
      55, 126, 184,
      77, 175, 74,
      255, 231, 255,
      229, 6, 255,
      96, 70, 0,
      88, 1, 132,
      255, 255, 0,
      0, 255, 245,
      166, 135, 135,
      116, 0, 42,
      255, 153, 0,
      0, 80, 73,
      0, 0, 229,
      179, 160, 255,
      136, 60, 255,
      195, 221, 172,
      255, 102, 162,
      0, 108, 0,
      97, 79, 92,
      24, 174, 167,
      0, 252, 19,
      154, 127, 2,
      189, 0, 128,
      116, 202, 255,
      94, 123, 100,
      255, 175, 162,
      155, 79, 58,
      189, 188, 0,
      255, 162, 250,
      152, 101, 164,
      181, 186, 194,
      47, 67, 105,
      33, 231, 153,
      221, 116, 79,
      219, 255, 255,
      156, 165, 117,
      149, 3, 167,
      137, 141, 171,
      159, 2, 0,
      255, 213, 105,
      95, 32, 0,
      54, 82, 255,
      0, 121, 127,
      98, 93, 154,
      116, 128, 255,
      185, 95, 126,
      106, 39, 91,
      90, 3, 211,
      42, 68, 1,
      190, 145, 191,
      138, 207, 192,
      182, 255, 129,
      45, 0, 137,
      238, 0, 105,
      75, 83, 64,
      182, 103, 252,
      255, 225, 200,
      127, 105, 63,
      99, 110, 122,
      206, 205, 255,
      96, 134, 0,
      26, 164, 255,
      244, 2, 187,
      130, 154, 150,
      181, 98, 0,
      4, 154, 108,
      211, 159, 105,
      146, 55, 85,
      161, 255, 212,
      144, 208, 107,
      223, 184, 207,
      231, 111, 212,
      0, 108, 71,
      107, 165, 193,
      106, 66, 50,
      88, 101, 29,
      254, 255, 181,
      186, 0, 63,
      218, 141, 158,
      174, 1, 242,
      0, 91, 165,
      0, 212, 226,
      255, 84, 94,
      115, 71, 131,
      0, 207, 0,
      204, 189, 167,
      111, 179, 138,
      255, 81, 0,
      139, 99, 106,
      157, 234, 255,
      127, 0, 202,
      200, 155, 26,
      186, 59, 178,
      143, 161, 45,
      165, 225, 0,
      173, 126, 81,
      0, 91, 114,
      158, 177, 221,
      186, 61, 0,
      67, 79, 86,
      199, 191, 111,
      216, 139, 255,
      138, 119, 197,
      201, 80, 77,
      129, 60, 0,
      122, 136, 75,
      80, 60, 129,
      216, 220, 218,
      131, 45, 38,
      125, 120, 114,
      1, 138, 53,
      3, 146, 172,
      146, 0, 116,
      217, 226, 84,
      132, 116, 138,
      150, 73, 129,
      222, 121, 3,
      255, 148, 94,
      211, 175, 235,
      102, 45, 58,
      244, 255, 233,
      75, 142, 133,
      16, 124, 255,
      255, 146, 200,
      255, 186, 122,
      163, 185, 167,
      0, 204, 165,
      65, 102, 95,
      133, 80, 184,
      238, 183, 0,
      118, 150, 213,
      255, 131, 128,
      146, 193, 209,
      221, 233, 255,
      1, 79, 43,
      49, 61, 158,
      117, 0, 0,
      178, 109, 100,
      140, 89, 0,
      168, 157, 171,
      133, 189, 13,
      82, 255, 129,
      105, 137, 152,
      189, 111, 170,
      102, 100, 215,
      252, 200, 255,
      71, 93, 125,
      108, 93, 83,
      94, 0, 59,
      201, 159, 149,
      142, 212, 160,
      85, 57, 95,
      214, 63, 139,
      193, 97, 211,
      104, 119, 160,
      162, 159, 148,
      220, 108, 122,
      0, 198, 119,
      86, 0, 93,
      114, 150, 117,
      0, 46, 123,
      0, 158, 0,
      201, 255, 187,
      174, 136, 215,
      176, 211, 247,
      186, 229, 213,
      0, 188, 229,
      145, 138, 107,
      83, 235, 199,
      67, 9, 255,
      255, 209, 217,
      106, 94, 123,
      230, 213, 155,
      213, 89, 17,
      168, 146, 73,
      104, 69, 192,
      172, 52, 46,
      158, 129, 255,
      71, 119, 64,
      113, 25, 129,
      184, 59, 100,
      216, 207, 226,
      120, 67, 97,
      148, 164, 179,
      149, 3, 41,
      124, 106, 3,
      174, 161, 205,
      163, 116, 140,
      62, 222, 99,
      122, 177, 255,
      133, 179, 103,
      142, 0, 77,
      78, 75, 104,
      252, 255, 123,
      68, 116, 141,
      48, 95, 39,
      199, 255, 0,
      121, 178, 175,
      164, 255, 246,
      253, 251, 255,
      89, 150, 81,
      0, 127, 105,
      139, 138, 144,
      131, 74, 77,
      255, 49, 236,
      255, 215, 0,
      95, 104, 70,
      73, 79, 28,
      71, 84, 183,
      186, 214, 222,
      185, 126, 31,
      74, 14, 153,
      255, 0, 70,
      156, 130, 172,
      157, 100, 54,
      167, 193, 139,
      230, 99, 255,
      152, 238, 154,
      204, 133, 119,
      148, 113, 94,
      101, 84, 48,
      255, 119, 39,
      255, 174, 197,
      223, 123, 176,
      255, 45, 151,
      152, 69, 169,
      131, 93, 124,
      99, 50, 255,
      129, 81, 44,
      194, 5, 0,
      159, 76, 7,
      227, 176, 83,
      92, 2, 14,
      142, 154, 255,
      163, 81, 92,
      97, 105, 100,
      213, 81, 182,
      113, 130, 128,
      225, 160, 210,
      186, 94, 58,
      103, 162, 0,
      216, 236, 147,
      96, 54, 12,
      0, 177, 140,
      61, 151, 203,
      215, 38, 70,
      153, 50, 6,
      98, 126, 203,
      199, 159, 178,
      231, 79, 55,
      40, 107, 203,
      173, 1, 153,
      125, 106, 157,
      255, 84, 195], dtype=jnp.uint8).reshape(256, 3)
